from conf.config import settings
from openai import OpenAI
from .base_llm import BaseLLM
from tiktoken import encoding_for_model
import transformers
from llmadapter.models import DEEPSEEK_TOKENIZER_DIR, QWEN_TOKENIZER_DIR

class LLMOpenAI(BaseLLM):
    def __init__(self, model: str, api_key: str, base_url: str):
        super().__init__()
        self.client = OpenAI(
            base_url=base_url,
            api_key=api_key
        )
        self.model = model
        self.encoding = None

    def _encoding(self):
        if self.encoding is not None:
            return self.encoding

        try:
            self.encoding = encoding_for_model(self.model)
        except Exception as e:
            if self.model.startswith("deepseek"):
                self.encoding = transformers.AutoTokenizer.from_pretrained(DEEPSEEK_TOKENIZER_DIR, trust_remote_code=True)
            elif self.model.startswith("qwen"):
                self.encoding = transformers.AutoTokenizer.from_pretrained(QWEN_TOKENIZER_DIR, trust_remote_code=True)
            else:
                raise ValueError(f"Unsupported model: {self.model}")
            
        return self.encoding

    def get_max_input_tokens(self):
        return settings.LLM_CONFIG["MAX_INPUT_TOKENS"]

    def get_max_output_tokens(self):
        return settings.LLM_CONFIG["MAX_OUTPUT_TOKENS"]

    def get_context_token(self, text: str):
        # self.encoding = self._encoding()
        self.encoding = None
        return super().get_context_token(self.model, text)

    def tokenize_encode(self, text: str):
        try:
            self.encoding = self._encoding()
            return self.encoding.encode(text)
        except Exception as e:
            raise ValueError(f"Failed to tokenize encode: {e}")

    def tokenize_decode(self, tokens: list):
        try:
            self.encoding = self._encoding()
            return self.encoding.decode(tokens)
        except Exception as e:
            raise ValueError(f"Failed to tokenize decode: {e}")

    def predict(self, messages: list, **kwargs):
        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                **kwargs
            )
            return response.choices[0].message.content
        except Exception as e:
            raise ValueError(f"Failed to predict: {e}")

    async def apredict(self, messages: list, **kwargs):
        try:
            response = await self.client.chat.completions.acreate(
                model=self.model,
                messages=messages,
                **kwargs
            )
            return response.choices[0].message.content
        except Exception as e:
            raise ValueError(f"Failed to predict asynchronously: {e}")
